# 行分割保存到excel里且支持追加保存
# 支持保存为txt格式
# 增加“帮助”按钮
# 优化界面
from PyQt6.QtWidgets import (QApplication, QInputDialog, QFileDialog,
                             QProgressDialog)
from PyQt6.QtGui import QPixmap, QDragEnterEvent, QDropEvent
from PyQt6.QtCore import Qt, QSize
import cv2
from paddleocr import PaddleOCR
import pandas as pd
import os
from datetime import datetime
from io import StringIO
import sys
from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QLabel,
                             QPushButton, QTextEdit, QComboBox, QMessageBox)

# try:
#     import pyi_splash
#     pyi_splash.close()
# except ImportError:
#     pass

class BaseApp(QWidget):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.setMinimumSize(QSize(400, 600))  # 设置最小窗口尺寸

    def init_ui(self):
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(10, 10, 10, 10)  # 减少边距
        main_layout.setSpacing(8)  # 减小控件间距

        # 上传区域和识别语言区域
        top_layout = QHBoxLayout()
        top_layout.setSpacing(8)  # 减小控件间距

        # 上传区域
        upload_layout = QHBoxLayout()
        upload_layout.setSpacing(8)  # 减小控件间距
        self.upload_btn = QPushButton("📁 上传图片")
        self.upload_btn.setStyleSheet("font-size: 14px; padding: 6px;")
        self.upload_btn.setFixedWidth(100)  # 设置上传按钮宽度
        self.help_btn = QPushButton("❓ 帮助")  # 添加帮助按钮
        self.help_btn.setStyleSheet("font-size: 14px; padding: 6px;")
        self.help_btn.setFixedWidth(80)  # 设置帮助按钮宽度
        upload_layout.addWidget(self.upload_btn)
        upload_layout.addWidget(self.help_btn)  # 帮助按钮放在上传按钮右边
        top_layout.addLayout(upload_layout)

        # 识别语言区域
        lang_layout = QHBoxLayout()
        lang_layout.setSpacing(4)  # 减小控件间距
        lang_layout.addWidget(QLabel("识别语言:"))
        self.lang_combo = QComboBox()
        self.lang_combo.addItems(["中文", "英文"])
        self.lang_combo.setFixedWidth(100)
        lang_layout.addWidget(self.lang_combo)
        lang_layout.addStretch()  # 添加伸缩空间，使语言选择靠左
        top_layout.addLayout(lang_layout)

        main_layout.addLayout(top_layout)

        # 图片预览
        self.image_label = QLabel()
        self.image_label.setFixedSize(320, 240)  # 缩小预览区域
        self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.image_label.setStyleSheet("""border: 2px dashed #ccc;
            background-color: #f8f9fa;
            border-radius: 8px;""")
        main_layout.addWidget(self.image_label)

        # OCR按钮
        self.ocr_btn = QPushButton("🔍 开始识别")
        self.ocr_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                border: none;
                padding: 10px;
                border-radius: 6px;
                font-size: 14px;
            }
            QPushButton:hover { background-color: #45a049; }
        """)
        main_layout.addWidget(self.ocr_btn)

        # 识别结果
        self.ocr_result = QTextEdit()
        self.ocr_result.setPlaceholderText("识别结果将显示在此处...")
        self.ocr_result.setStyleSheet("""
            QTextEdit {
                border: 1px solid #ddd;
                border-radius: 6px;
                padding: 8px;
                font-size: 13px;
            }
        """)
        main_layout.addWidget(self.ocr_result)

        # 底部操作栏
        bottom_layout = QHBoxLayout()
        bottom_layout.setSpacing(8)  # 减小控件间距
        self.select_folder_btn = QPushButton("📂 选择文件夹")
        self.select_folder_btn.setStyleSheet("font-size: 13px; padding: 6px;")
        self.select_folder_btn.setFixedWidth(120)  # 设置选择文件夹按钮宽度
        self.input_filename_btn = QPushButton("✏️ 输入文件名")
        self.input_filename_btn.setStyleSheet("font-size: 13px; padding: 6px;")
        self.input_filename_btn.setFixedWidth(120)  # 设置输入文件名按钮宽度
        self.save_excel_btn = QPushButton("💾 保存结果")
        self.save_excel_btn.setStyleSheet("font-size: 13px; padding: 6px;")
        self.save_excel_btn.setFixedWidth(120)  # 设置保存结果按钮宽度

        bottom_layout.addWidget(self.select_folder_btn)
        bottom_layout.addWidget(self.input_filename_btn)
        bottom_layout.addWidget(self.save_excel_btn)
        bottom_layout.addStretch()  # 添加伸缩空间，使按钮靠左
        main_layout.addLayout(bottom_layout)

        self.setLayout(main_layout)

    def show_help(self):
        help_text = """
        <h3>OCR文字识别工具 - 使用指南</h3>
        <p><strong>上传图片</strong>: 点击“上传图片”按钮选择要识别的图片文件。</p>
        <p><strong>选择语言</strong>: 在“识别语言”下拉菜单中选择图片文字的语言（中文或英文）。</p>
        <p><strong>开始识别</strong>: 点击“开始识别”按钮对上传的图片进行文字识别。</p>
        <p><strong>查看结果</strong>: 识别结果会显示在下方的文本框中。</p>
        <p><strong>选择文件夹</strong>: 点击“选择文件夹”按钮选择保存结果的文件夹。</p>
        <p><strong>输入文件名</strong>: 点击“输入文件名”按钮为保存的文件命名。</p>
        <p><strong>保存结果</strong>: 点击“保存结果”按钮将识别结果保存为Excel或文本文件。</p>
        """
        QMessageBox.about(self, "帮助", help_text)

class MyApp(BaseApp):
    def __init__(self):
        super().__init__()
        self.ocr_engine = None
        self.current_image_path = None
        self.output_folder = os.path.expanduser("~/Desktop")  # 默认保存到桌面
        self.excel_file_name = "OCR_Results.xlsx"
        self.processing = False  # 防止重复点击

        # 初始化OCR引擎
        self.init_ocr_engine()

        # 连接信号
        self.upload_btn.clicked.connect(self.open_image_dialog)
        self.ocr_btn.clicked.connect(self.run_ocr)
        self.select_folder_btn.clicked.connect(self.select_output_folder)
        self.input_filename_btn.clicked.connect(self.input_file_name)
        self.save_excel_btn.clicked.connect(self.save_results)
        self.lang_combo.currentIndexChanged.connect(self.init_ocr_engine)
        self.help_btn.clicked.connect(self.show_help)

        self.setAcceptDrops(True)
        self.setWindowTitle("OCR文字识别工具")

    def init_ocr_engine(self):
        lang = 'ch' if self.lang_combo.currentText() == "中文" else 'en'
        try:
            self.ocr_engine = PaddleOCR(use_angle_cls=True, lang=lang)
        except Exception as e:
            QMessageBox.critical(self, "引擎初始化失败", f"无法初始化OCR引擎：{str(e)}")

    def open_image_dialog(self):
        if self.processing:
            return
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if file_path:
            self.display_image(file_path)

    def display_image(self, path):
        self.current_image_path = path
        pixmap = QPixmap(path)
        if not pixmap.isNull():
            scaled_pixmap = pixmap.scaled(
                self.image_label.width(),
                self.image_label.height(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.image_label.setPixmap(scaled_pixmap)
        else:
            self.show_error("图片加载失败", "不支持的文件格式或损坏的图片文件")

    def preprocess_image(self, image_path):
        try:
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError("无法读取图片文件")

            # 显示预处理进度
            progress = QProgressDialog("正在预处理图片...", None, 0, 0, self)
            progress.setWindowModality(Qt.WindowModality.WindowModal)
            progress.show()

            QApplication.processEvents()  # 更新UI

            # 预处理流程
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            binary = cv2.adaptiveThreshold(
                gray, 255,
                cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                cv2.THRESH_BINARY, 11, 2
            )
            temp_path = os.path.join(os.path.dirname(__file__), "temp_preprocessed.png")
            cv2.imwrite(temp_path, binary)

            progress.close()
            return temp_path

        except Exception as e:
            self.show_error("预处理失败", str(e))
            return None

    def run_ocr(self):
        if self.processing:
            return

        if not self.current_image_path:
            self.show_warning("提示", "请先选择要识别的图片")
            return

        try:
            self.processing = True
            self.ocr_result.setPlainText("识别中，请稍候...")
            QApplication.processEvents()  # 立即更新UI

            # 预处理和OCR识别
            processed_path = self.preprocess_image(self.current_image_path)
            if processed_path is None:
                return

            result = self.ocr_engine.ocr(processed_path, cls=True)

            # 解析结果
            texts = [line[1][0] for line in result[0] if len(line) > 1]
            self.ocr_result.setPlainText("\n".join(texts))

            # 自动滚动到最后
            cursor = self.ocr_result.textCursor()
            cursor.movePosition(cursor.MoveOperation.End)
            self.ocr_result.setTextCursor(cursor)

        except Exception as e:
            self.show_error("识别失败", f"发生错误：{str(e)}")
        finally:
            self.processing = False
            if os.path.exists(processed_path):
                os.remove(processed_path)  # 清理临时文件

    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QDropEvent):
        if self.processing:
            return
        for url in event.mimeData().urls():
            file_path = url.toLocalFile()
            if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                self.display_image(file_path)
                break

    def select_output_folder(self):
        folder_path = QFileDialog.getExistingDirectory(self, "选择保存位置", self.output_folder)
        if folder_path:
            self.output_folder = folder_path
            QMessageBox.information(self, "路径已更新", f"文件将保存到：\n{folder_path}")

    def input_file_name(self):
        text, ok = QInputDialog.getText(
            self, "输入文件名",
            "请输入文件名（不包含扩展名）：",
            text=self.excel_file_name.split(".")[0]  # 默认文件名（去掉扩展名）
        )
        # if ok and text:
        #     # 弹出对话框让用户选择保存格式
        #     format, format_ok = QInputDialog.getItem(
        #         self, "选择保存格式",
        #         "请选择保存格式：",
        #         ["Excel (.xlsx)", "Text (.txt)"],
        #         0, False
        #     )
        #     if format_ok:
        #         if format == "Excel (.xlsx)":
        #             self.excel_file_name = f"{text}.xlsx"
        #         elif format == "Text (.txt)":
        #             self.excel_file_name = f"{text}.txt"

    def save_results(self):
        # 弹出对话框让用户选择保存格式
        format, ok = QInputDialog.getItem(
            self, "选择保存格式",
            "请选择保存格式：",
            ["Excel (.xlsx)", "Text (.txt)"],
            0, False
        )
        if not ok:
            return  # 用户取消选择

        if format == "Excel (.xlsx)":
            self.save_to_excel()
        elif format == "Text (.txt)":
            self.save_to_text()

    def save_to_text(self):
        if not self.ocr_result.toPlainText().strip():
            self.show_warning("保存失败", "没有可保存的识别结果")
            return

        try:
            text_path = os.path.join(self.output_folder, self.excel_file_name.replace(".xlsx", ".txt"))
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            # 将识别结果按行分割
            ocr_lines = self.ocr_result.toPlainText().split('\n')

            # 写入文本文件
            with open(text_path, "a", encoding="utf-8") as file:
                file.write(f"时间: {timestamp}\n")
                file.write(f"图片路径: {self.current_image_path}\n")
                file.write("识别结果:\n")
                for line in ocr_lines:
                    file.write(f"{line}\n")
                file.write("\n")

            self.show_info("保存成功", f"文件已保存到：\n{text_path}")

        except Exception as e:
            self.show_error("保存失败", f"写入文本文件时出错：\n{str(e)}")

    def save_to_excel(self):
        # 添加依赖检查
        try:
            import openpyxl
        except ImportError as e:
            self.show_error(
                "缺失关键依赖库",
                f"无法执行保存操作，原因：{str(e)}\n\n"
                "请通过以下命令安装所需库：\n"
                "pip install openpyxl"
            )
            return

        # 原有保存逻辑
        if not self.ocr_result.toPlainText().strip():
            self.show_warning("保存失败", "没有可保存的识别结果")
            return

        try:
            excel_path = os.path.join(self.output_folder, self.excel_file_name)
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            # 将识别结果按行分割
            ocr_lines = self.ocr_result.toPlainText().split('\n')

            # 创建数据框，每行对应OCR的一行
            df = pd.DataFrame({
                "时间": [timestamp] * len(ocr_lines),
                "图片路径": [self.current_image_path] * len(ocr_lines),
                "识别结果": ocr_lines
            })

            if os.path.exists(excel_path):
                # 如果文件存在，追加到现有工作表
                with pd.ExcelWriter(
                        excel_path,
                        mode='a',
                        engine='openpyxl',
                        if_sheet_exists='overlay'  # 覆盖模式，追加到现有工作表
                ) as writer:
                    # 检查工作表是否存在
                    if '结果' in writer.book.sheetnames:
                        # 如果存在，读取现有数据并追加
                        existing_df = pd.read_excel(excel_path, sheet_name='结果')
                        combined_df = pd.concat([existing_df, df], ignore_index=True)
                        combined_df.to_excel(
                            writer,
                            index=False,
                            sheet_name='结果'
                        )
                    else:
                        # 如果不存在，创建新工作表
                        df.to_excel(
                            writer,
                            index=False,
                            sheet_name='结果'
                        )
            else:
                # 如果文件不存在，创建新文件和工作表
                df.to_excel(
                    excel_path,
                    index=False,
                    sheet_name="结果",
                    engine='openpyxl'  # 显式指定引擎
                )

            self.show_info("保存成功", f"文件已保存到：\n{excel_path}")

        except Exception as e:
            self.show_error("保存失败", f"写入Excel文件时出错：\n{str(e)}")

    # 辅助提示方法
    def show_error(self, title, message):
        QMessageBox.critical(self, title, message)

    def show_warning(self, title, message):
        QMessageBox.warning(self, title, message)

    def show_info(self, title, message):
        QMessageBox.information(self, title, message)

    def show_help(self):
        help_text = """
        <h3>OCR文字识别工具 - 使用指南</h3>
        <p><strong>上传图片</strong>: 点击“上传图片”按钮选择要识别的图片文件（支持拖动）。</p>
        <p><strong>选择语言</strong>: 在“识别语言”下拉菜单中选择图片文字的语言（中文或英文）。</p>
        <p><strong>开始识别</strong>: 点击“开始识别”按钮对上传的图片进行文字识别。</p>
        <p><strong>查看结果</strong>: 识别结果会显示在下方的文本框中。</p>
        <p><strong>选择文件夹</strong>: 点击“选择文件夹”按钮选择保存结果的文件夹。</p>
        <p><strong>输入文件名</strong>: 点击“输入文件名”按钮为保存的文件命名。</p>
        <p><strong>保存结果</strong>: 点击“保存结果”按钮将识别结果保存为Excel或文本文件。</p>
        """
        QMessageBox.about(self, "帮助", help_text)

    if sys.stderr is None:
        sys.stderr = StringIO()

if __name__ == "__main__":
    import sys
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 使用更现代的主题
    window = MyApp()
    window.show()
    sys.exit(app.exec())
